# -*- coding: utf-8 -*-
# @Time    : 2018/7/21 21:14
# @Author  : Tianchiyue
# @File    : keyword_extraction.py
# @Software: PyCharm Community Edition

import jieba.analyse
from rake_nltk import Rake


def get_keywords(sentence_list, mode='rake', language='en'):
    """
    :param sentence_list:
    :param mode:tfidf, textrank, rake
    :param language: english, chinese
    :return:
    """
    text = ' '.join(sentence_list)
    r = Rake()
    result = []
    if mode == 'rake' and language == 'en':
        r.extract_keywords_from_text(text)
        result = r.get_ranked_phrases()[:10]
    elif mode == 'tfidf' and language == 'cn':
        result = jieba.analyse.extract_tags(text, topK=10)
    elif mode == 'textrank' and language == 'cn':
        result = jieba.analyse.textrank(text, topK=10)
    return result
